from matplotlib import pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs


def test1():
    points, cluster_indexes = make_blobs(n_samples=300, centers=4, cluster_std=0.8, random_state=0)

    x = points[:, 0]
    y = points[:, 1]

    plt.scatter(x, y, s=50, alpha=0.7)

    plt.show()


def test2():
    points, cluster_indexes = make_blobs(n_samples=300, centers=4, cluster_std=0.8, random_state=0)

    x = points[:, 0]
    y = points[:, 1]

    plt.scatter(x, y, alpha=0.7)
    plt.show()

    kmeans = KMeans(n_clusters=4, random_state=0)
    kmeans.fit(points)
    predicted_cluster_indexes = kmeans.predict(points)

    plt.scatter(x, y, c=predicted_cluster_indexes, s=50, alpha=0.7)
    # plt.show()
    centers = kmeans.cluster_centers_
    plt.scatter(centers[:, 0], centers[:, 1], color='red', s=100)

    plt.show()
